import numpy as np 
import numpy.random as npr
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.optimize import leastsq,least_squares
from sklearn import preprocessing

def logistic5(x,p1,p2,p3,p4,p5):
    c = 2*p3*p5/abs(p3+p5)
    f = 1/(1 + np.exp(-c*(p4-x)))
    g = np.exp(p3*(p4-x))
    h = np.exp(p5*(p4-x))
    return p1 + (p2/(1 + f*g + (1-f)*h))

def residuals5(p,y,x):
    p1,p2,p3,p4,p5 = p
    err = y - logistic5(x,p1,p2,p3,p4,p5)
    return err

def peval5(x,p):
    p1,p2,p3,p4,p5 = p
    return logistic5(x,p1,p2,p3,p4,p5)
'''
x = np.linspace(0,20,20)
p1,p2,p3,p4,p5 = 0.5,2.5,8,7.3,3.9
y_true = logistic5(x,p1,p2,p3,p4,p5)
y_meas = y_true + 0.2*npr.randn(len(x))
'''
def get_data():
    Work_Path = '/home/lss/Lab/super_com/VASP/old_file/'
    #row_log = open(Work_Path + 'row_log_normal_no_converg.csv')
    row_log = open(Work_Path + 'row_log_normal_converg.csv')
    line = row_log.readline().strip().split(',')
    conv_steps = {}
    
    while not line == ['']:
        id_step = []
        for i in range(1,len(line)):
            if line[i]:
                id_step.append(float(line[i]))
        #print id_step
        conv_steps[line[0]] = id_step
        line = row_log.readline().strip().split(',')
    return conv_steps

if __name__ == '__main__':
    conv_steps = get_data()
    #print conv_steps
    p0 = [0,0,10,0,0]
    min_max_scaler = preprocessing.MinMaxScaler()
    for job_id in conv_steps:
        x = []
        y_true = [] 
        for i in range(1,21):
            x.append(i)
        x = np.array(x)
        #print len(conv_steps[job_id]) 
        #if len(conv_steps[job_id]) >= 50:
        if len(conv_steps[job_id]) >= 40:
            #print conv_steps[job_id]
            y_true = np.array(conv_steps[job_id][15:35])            
            y_scaler = min_max_scaler.fit_transform(y_true)
            #print y_true
            plt.figure()
            plsq = least_squares(residuals5,p0,f_scale=0.1,args=(y_scaler,x))
            #plt.plot(x,peval5(x,plsq.x),x,y_true)
            x_plot = []
            for i in range(1,26):
                x_plot.append(i)
            x_plot = np.array(x_plot)
            y_plot = np.array(conv_steps[job_id][15:40])
            
            plt.plot(x_plot,peval5(x_plot,plsq.x),x_plot,y_plot)
            plt.title(job_id+'(scaler predict 5 pots) Least-squares 5PL ')
            plt.legend(['Fit','True'],loc = 'upper left')
            for i , (param, est) in enumerate(zip('ABCDE',plsq.x)):
                plt.text(10,3-i*0.5,'est(%s) = %.2f' % (param, est))
            plt.savefig('logis5_' + job_id +'.png')
        
    '''
    # Initial guess for parameters
    p0 = [0,0,10,0,0]

    # Fit equation using least squares optimization
    #plsq = leastsq(residuals,p0,args=(y_meas,x))
    plsq = least_squares(residuals5,p0,f_scale=0.1,args=(y_meas,x))
    #plsq = least_squares(residuals5,p0,f_scale=0.1,args=(x,y_meas))

    plt.plot(x,peval5(x,plsq.x),x,y_meas,'o',x,y_true)
    plt.title('Least-squares 5PL fit noisy data')
    plt.legend(['Fit','Noisy','True'],loc = 'upper left')
    for i , (param, actual, est) in enumerate(zip('ABCDE',[p1,p2,p3,p4,p5],plsq.x)):
        plt.text(10,3-i*0.5,'%s = %.2f, est(%s) = %.2f' % (param, actual, param, est))
    plt.savefig('logis_text5_new.png')
    #plt.show()
    '''
